# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： activation
Description :
Author : 'li'
date： 2022/7/3
Change Activity:
2022/7/3:
-------------------------------------------------
"""

from torch import nn

from ....common.design_pattern.register.register_factory import REGISTER_FACTORY

_registry_name = 'activation'


def _build_registry():
    """
    build normal lay registry.

    """
    _register = REGISTER_FACTORY.create_registry(_registry_name)
    _register.add('ReLU', value=nn.ReLU)
    _register.add('LeakyReLU', value=nn.LeakyReLU)
    _register.add('PReLU', value=nn.PReLU)
    _register.add('RReLU', value=nn.RReLU)
    _register.add('ReLU6', value=nn.ReLU6)
    _register.add('ELU', value=nn.ELU)
    _register.add('Sigmoid', value=nn.Sigmoid)
    _register.add('Tanh', value=nn.Tanh)


_build_registry()  # build when module load.


def build_activation_layer(cfg=None, name=None):
    """
    build norm layer

    Args:
        name: module name
        cfg: must contain num_features;

    Returns:

    """
    if isinstance(cfg, str):
        cfg = dict(type=cfg)
    _type = cfg['type']
    if _type in ['ReLU']:
        cfg['inplace'] = True
    return REGISTER_FACTORY.build_from_config(_registry_name, cfg, name)
